##__docformat__ = "restructuredtext en"
# ******NOTICE***************
# two_opt.py module by Shoibal Chakravarty
#
# You may copy and use this module as you see fit with no
# guarantee implied provided you keep this notice in all copies.
# *****END NOTICE************
"""
This module  contains the functions that simulate a scenario where two 
regions have a two period planning horizon (2nd period myopic), and the third is myopic. The two 
regions cooperate to maximize the sum of their two period utility functions.
The green energy deployment in the last period is myopic.
"""

__author__ = "Shoibal Chakravarty (shoibalc@princeton.edu)"
__version__ = "0.1"
__date__ = "2011/August"
__license__ = "BSD"
__all__ = ['twooptimized','twooptimizedupdate','twooptimizedregion']
import numpy as np
from myopic import  regionupdate, worldupdate, myopicupdate, cumulativecost
from flat2 import flatten
from utility import U2, U2optimized
from two_game import finalperiodupdate
import scipy.optimize

def twooptimized(r1,r2,r3,r4,w,com,precision=0.5,init_grid_size=40):
    """
    Function twooptimized simulates a scenario where two regions have a two 
    period planning horizon (current period and the next period, where each region is myopic), 
    and the third region is myopic. The two regions cooperate to maximize the sum of their two 
    period utility functions. The solution involves solving a non-linear optimization problem to 
    determine the best current period deployments of these two regions. The green energy deployment
    in the last period is myopic.\n\n
    twooptimized(r1,r2,r3,...):\n
    r1 and r2 solve the two region two period optimial green energy deployment problem, 
    r3 is myopic.
    """
    r1.descriptor = r2.descriptor = 'Optimal solution for two regions with one period look ahead.'
    r3.descriptor = r4.descriptor = 'Myopic'
    print "\n\nIn this scenario %s and %s  play a two region two period Nash game, %s and %s are myopic\n"% (r1.name, r2.name, r3.name, r4.name)
#  handle the myopic region r3 first
    for t in range(w.T-1):
        r3.g[t], r3.u1[t], r3.c1[t], r3.m[t] = myopicupdate(r3,r1,r2,r4,w,t)
        regionupdate(r3,t)
        r4.g[t], r4.u1[t], r4.c1[t], r4.m[t] = myopicupdate(r4,r1,r2,r3,w,t)
        regionupdate(r4,t)
        twooptimizedupdate(r1,r2,r3,r4,w,com,t,precision,init_grid_size)
        worldupdate(r1,r2,r3,r4,w,com,t)
        print "\nPeriod", t+1,"\n", r1.name, "g", r1.g[t], "e", r1.e[t],"\n", r2.name, "g", r2.g[t], "e", r2.e[t],\
        "\n",r3.name, "g", r3.g[t], "e", r3.e[t], "\n",r4.name, "g", r4.g[t], "e", r4.e[t], "\ncg", w.cg[t+1],"\n"
# the last period is still myopic
    t = w.T-1
    finalperiodupdate(r1,t)
    finalperiodupdate(r2,t)
    r3.g[t], r3.u1[t], r3.c1[t], r3.m[t] = myopicupdate(r3,r1,r2,r4,w,t)
    regionupdate(r3,t)
    r4.g[t], r4.u1[t], r4.c1[t], r4.m[t] = myopicupdate(r4,r1,r2,r3,w,t)
    regionupdate(r4,t)
    worldupdate(r1,r2,r3,r4,w,com,t)
    print "\nPeriod", t+1,"\n", r1.name, "g", r1.g[t], "e", r1.e[t],"\n", r2.name, "g", r2.g[t], "e", r2.e[t],\
    "\n",r3.name, "g", r3.g[t], "e", r3.e[t], "\n",r4.name, "g", r4.g[t], "e", r4.e[t], "\ncg", w.cg[t+1],"\n"
    cumulativecost(r1,r2,r3,r4,w)

def twooptimizedupdate(ra,rb,rc,rd,w,com,t,precision,init_grid_size):
    """
    Function twooptimizedupdate solves the two region two period optimization problem 
    for the current period (except for the last period in the simulation). 
    The function successively calls function twooptimizedregion which return 
    the current best guess and a rectangle bounding the solution. 
    This procedures continues till the desired precision is reached.
    """
    maxa = min(ra.M[t]*ra.E[t+1],(ra.E[t+1] - ra.G[t]))
    mina = 1e-6
    maxb = min(rb.M[t]*rb.E[t+1],(rb.E[t+1] - rb.G[t]))
    minb = 1e-6
    D0 = 19
    D = D0
    EPS = EPSa = EPSb = init_grid_size*D
    print 'Search region',ra.name,'min = 0 max =', maxa, rb.name, 'min = 0, max =', maxb
    flag = True
    num = False
    xbest, ybest = np.nan, np.nan
    while EPS > precision and flag:
# initial run with  init_grid_size. Subsequent runs reduce it by a factor of D
        if num:
            D = np.random.randint(D0+1,D0+5)
        EPSa = EPSa/D
        EPSb = EPSb/D
        num = True
        print "GRID SIZE", D
        mina, maxa, minb, maxb, xbest1, ybest1, lenxy = \
                twooptimizedregion(ra,rb,rc,rd,w,com,t,mina,maxa,minb,maxb, EPSa,EPSb)

        print "Search region update:", ra.name, "min", mina, "max", maxa, "best", xbest1, rb.name, "min", minb, "max", maxb, "best", ybest1
        err_best = (xbest - xbest1)**2 + (ybest - ybest1)**2
#        if np.sqrt(err_best) <= 0.25*precision:
# this hack prevents infinite looping if we hit a multiple equilibria. we select a particular solution and carry on.
# it also helps in avoiding an extra round of brute force maximization in some cases.
#            if lenxy == 1:
#            print "1. ATTAINED REQUIRED PRECISION"
#            flag = False
#            else:
#                print "2. MULTIPLE SOLUTIONS DETECTED, PICKING AVERAGE OF DEGENERATE SOLUTIONS"
#                flag = False
#            break
#        else:
        xbest, ybest = xbest1, ybest1
        EPSa = maxa - mina
        EPSb = maxb - minb
        EPS = max(EPSa, EPSb)
    if flag:
        print "3. ATTAINED REQUIRED PRECISION"
        ga = 0.5*(mina + maxa)
        gb = 0.5*(minb + maxb)
    else:
        ga = xbest
        gb = ybest
    
    parama = \
    (gb, com.cg0, com.greenpar, ra.B[t], ra.B[t+1], ra.cn[t], ra.co[t], ra.cn[t+1], ra.co[t+1], ra.G[t], rb.G[t], \
    rc.G[t]+rd.G[t], rc.G[t+1]+rd.G[t+1], w.G[0], ra.e[t], ra.E[t+1], ra.e[t+1], ra.E[t+2], ra.M[t], ra.M[t+1], (1-com.discount)**com.years,True )
    u, ra.g[t], ra.u1[t], ra.c1[t], ra.g2[t], ra.u2[t], ra.c2[t], ra.m[t] =  U2(ga,*parama)
    regionupdate(ra,t)
    
    paramb = \
    (ga, com.cg0, com.greenpar, rb.B[t], rb.B[t+1], rb.cn[t], rb.co[t], rb.cn[t+1], rb.co[t+1], rb.G[t], ra.G[t], \
    rc.G[t]+rd.G[t], rc.G[t+1]+rd.G[t+1], w.G[0], rb.e[t], rb.E[t+1], rb.e[t+1], rb.E[t+2], rb.M[t], rb.M[t+1], (1-com.discount)**com.years,True )
    u, rb.g[t], rb.u1[t], rb.c1[t], rb.g2[t], rb.u2[t], rb.c2[t], rb.m[t] =  U2(gb,*paramb)
    regionupdate(rb,t)


def twooptimizedregion(ra,rb,rc,rd,w,com,t,mina,maxa,minb,maxb,EPSa,EPSb):
    """
    Function twooptimizedregion takes among its parameters the coordinates 
    of a rectangular search region and a grid-size for discretization. It uses a brute 
    force method (scipy.optimize.brute) to numerically obtain the best solution to the 
    optimization problem - maximum of the sum of the two region's two period utility 
    functions. 
    """
#   Ga0,Gb0,Gc0,Gc1,G0,ea0,ea1,Ea0,Ea1,eb0,eb1,Eb0,Eb1,delta,precision,alpha, mina, maxa, minb, maxb, EPSa, EPSb = args
    
    Wa =3
    lena = np.ceil((maxa - mina)/EPSa) + 1
    lenb = np.ceil((maxb - minb)/EPSb) + 1
    EPSa = (maxa - mina)/(lena-1)
    EPSb = (maxb - minb)/(lenb-1)
    eps = 1e-5
    slice_ab = np.s_[mina: maxa + eps: EPSa, minb: maxb + eps: EPSb,]
    mgrid_a = np.mgrid[slice_ab[0]]
    mgrid_b = np.mgrid[slice_ab[1]]
    
    parama = \
    [0.0, com.cg0, com.greenpar, ra.B[t], ra.B[t+1], ra.cn[t], ra.co[t], ra.cn[t+1], ra.co[t+1], ra.G[t], rb.G[t], \
    rc.G[t]+rd.G[t], rc.G[t+1]+rd.G[t+1], w.G[0], ra.e[t], ra.E[t+1], ra.e[t+1], ra.E[t+2], ra.M[t], ra.M[t+1], (1-com.discount)**com.years,False]
    
    paramb = \
    [0.0, com.cg0, com.greenpar, rb.B[t], rb.B[t+1], rb.cn[t], rb.co[t], rb.cn[t+1], rb.co[t+1], rb.G[t], ra.G[t], \
    rc.G[t]+rd.G[t], rc.G[t+1]+rd.G[t+1], w.G[0], rb.e[t], rb.E[t+1], rb.e[t+1], rb.E[t+2], rb.M[t], rb.M[t+1], (1-com.discount)**com.years,False]
    
    paramab = (parama,paramb)
    xbest, ybest = scipy.optimize.brute(U2optimized,slice_ab,args=paramab,full_output=False,finish=None)
    
    xa  = [int(round((xbest - mina)/EPSa))]
    xb  = [int(round((ybest - minb)/EPSb))]
    xa1 = [xx - Wa > 0 and mgrid_a[xx - Wa].tolist() or mgrid_a[0].tolist() for xx in xa]
    xb1 = [xx - Wa > 0 and mgrid_b[xx - Wa].tolist() or mgrid_b[0].tolist() for xx in xb]
    xa2 = [xx + Wa < (lena - 1) and mgrid_a[xx + Wa].tolist() or mgrid_a[-1].tolist() for xx in xa]
    xb2 = [xx + Wa < (lenb - 1) and mgrid_b[xx + Wa].tolist() or mgrid_b[-1].tolist() for xx in xb]
#    print xa1,'\n',xb1,'\n',xc1,'\n',xa2,'\n',xb2,'\n',xc2,'\n'
    mina, maxa, minb, maxb = min(flatten(xa1)),max(flatten(xa2)),min(flatten(xb1)),max(flatten(xb2))
#    print mina, maxa, minb, maxb
#    print "min a", mina, "max a", maxa, "min b", minb, "max b", maxb, "best a", xbest, "best b", ybest
    return mina, maxa, minb, maxb, xbest, ybest, 1


